import httpx
import json
import logging
from typing import Dict, Any, Optional
from fastapi import HTTPException
from fastapi.responses import StreamingResponse
from config import settings

logger = logging.getLogger(__name__)

class OpenAIProxy:
    """OpenAI API代理类"""
    
    def __init__(self):
        # 使用httpx默认配置，只设置必要的超时
        self.client = httpx.AsyncClient(timeout=60.0)
    
    def _is_streaming_request(self, body: Optional[bytes]) -> bool:
        """检查是否为流式请求"""
        if not body:
            return False
        
        try:
            data = json.loads(body.decode('utf-8'))
            return data.get('stream', False)
        except (json.JSONDecodeError, UnicodeDecodeError):
            return False
    
    async def forward_request(
        self,
        method: str,
        path: str,
        headers: Dict[str, str],
        body: Optional[bytes] = None,
        backend: str = "openai"
    ) -> Dict[str, Any]:
        """转发请求到指定的后端服务"""
        # 检查是否为流式请求
        if self._is_streaming_request(body):
            return await self.forward_streaming_request(method, path, headers, body, backend)
        
        return await self._forward_regular_request(method, path, headers, body, backend)
    
    async def _forward_regular_request(
        self,
        method: str,
        path: str,
        headers: Dict[str, str],
        body: Optional[bytes],
        backend: str
    ) -> Dict[str, Any]:
        """转发普通请求"""
        try:
            # 获取后端配置
            backend_config = settings.backends.get(backend)
            if not backend_config:
                raise HTTPException(status_code=400, detail=f"不支持的后端服务: {backend}")
            
            # 构建URL和请求头
            url = f"{backend_config['base_url'].rstrip('/')}/{path.lstrip('/')}"
            proxy_headers = self._prepare_headers(headers)
            
            logger.info(f"转发请求到 {backend}: {method} {url}")
            
            # 发送请求
            response = await self.client.request(
                method=method,
                url=url,
                headers=proxy_headers,
                content=body
            )
            
            # 返回响应
            response_data = {
                "status_code": response.status_code,
                "headers": dict(response.headers),
                "content": response.content
            }
            
            # 尝试解析JSON响应
            if "application/json" in response.headers.get("content-type", ""):
                try:
                    response_data["json"] = response.json()
                except json.JSONDecodeError:
                    pass
            
            return response_data
            
        except httpx.RequestError as e:
            logger.error(f"请求错误: {e}")
            raise HTTPException(status_code=502, detail=f"代理请求失败: {str(e)}")
        except Exception as e:
            logger.error(f"未知错误: {e}")
            raise HTTPException(status_code=500, detail=f"内部服务器错误: {str(e)}")
    
    async def forward_streaming_request(
        self,
        method: str,
        path: str,
        headers: Dict[str, str],
        body: Optional[bytes] = None,
        backend: str = "openai"
    ) -> StreamingResponse:
        """转发流式请求"""
        try:
            # 获取后端配置
            backend_config = settings.backends.get(backend)
            if not backend_config:
                raise HTTPException(status_code=400, detail=f"不支持的后端服务: {backend}")
            
            # 构建URL和请求头
            url = f"{backend_config['base_url'].rstrip('/')}/{path.lstrip('/')}"
            proxy_headers = self._prepare_headers(headers)
            
            logger.info(f"转发流式请求到 {backend}: {method} {url}")
            
            return StreamingResponse(
                self._stream_generator(method, url, proxy_headers, body),
                media_type="text/plain"
            )
            
        except HTTPException:
            raise
        except Exception as e:
            logger.error(f"流式请求错误: {e}")
            raise HTTPException(status_code=500, detail=f"流式请求失败: {str(e)}")
    
    async def _stream_generator(
        self,
        method: str,
        url: str,
        headers: Dict[str, str],
        body: Optional[bytes]
    ):
        """流式响应生成器"""
        try:
            async with self.client.stream(
                method=method,
                url=url,
                headers=headers,
                content=body
            ) as response:
                if response.status_code >= 400:
                    error_content = await response.aread()
                    raise HTTPException(
                        status_code=response.status_code, 
                        detail=error_content.decode('utf-8', errors='ignore')
                    )
                
                async for chunk in response.aiter_bytes():
                    if chunk:
                        yield chunk
                        
        except HTTPException:
            raise
        except Exception as e:
            logger.error(f"流式生成器错误: {e}")
            raise HTTPException(status_code=500, detail=f"流式处理错误: {str(e)}")
    
    def _extract_api_key(self, headers: Dict[str, str]) -> Optional[str]:
        """从请求头中提取API密钥"""
        # 检查Authorization头
        auth_header = headers.get("authorization") or headers.get("Authorization")
        if auth_header and auth_header.startswith("Bearer "):
            return auth_header[7:]  # 移除 "Bearer " 前缀
        
        return None
    
    def _prepare_headers(self, original_headers: Dict[str, str]) -> Dict[str, str]:
        """准备转发的请求头"""
        # 提取API密钥
        api_key = self._extract_api_key(original_headers)
        if not api_key:
            raise HTTPException(status_code=401, detail="未提供API密钥")
        
        # 准备新的请求头
        headers = {
            "authorization": f"Bearer {api_key}",
            "content-type": original_headers.get("content-type", "application/json")
        }
        
        return headers
    
    async def close(self):
        """关闭HTTP客户端"""
        await self.client.aclose()

# 全局代理实例
proxy = OpenAIProxy() 